# coding: utf-8

from numpy import array, eye, zeros
from numpy import matmul, round, float64
from .util import transform_point_cloud


class Aligner(object):
    """
    ICP aligner. ICP has two major components:
    1. find correspondences;
    2. extract rotation;
    3. iteration.
    """

    def __init__(self, extract_rotation, find_correspondences, criteria, precision=10):
        super(Aligner, self).__init__()
        self._extract_rotation = extract_rotation
        self._find_correspondences = find_correspondences
        self._criteria = criteria
        self._precision = precision

    @property
    def extract_rotation(self):
        return self._extract_rotation

    @extract_rotation.setter
    def extract_rotation(self, extract_rotation):
        self._extract_rotation = extract_rotation

    @property
    def find_correspondences(self):
        return self._find_correspondences

    @find_correspondences.setter
    def find_correspondences(self, func):
        self._find_correspondences = func

    @property
    def stop_criteria(self):
        return self._criteria

    @stop_criteria.setter
    def stop_criteria(self, criteria):
        self._criteria = criteria

    def align(self, source_cloud, target_cloud, init_rot=eye(3)):
        translation_vector = zeros(3)
        rotation_matrix = init_rot
        should_stop = self._criteria()

        # initialize the clouds.
        # translated_cloud = array(source_cloud, dtype=float64, copy=False)
        # closest_cloud = array(target_cloud, dtype=float64, copy=False)

        # find the closest points.
        # translated_cloud, closest_cloud =self._find_correspondences(source_cloud, target_cloud)

        # delta_translation_vector = zeros(3)
        # delta_rotation_matrix = eye(3)

        max_dist = 4

        while True:
            # update the current transformation.
            translated_cloud = transform_point_cloud(source_cloud, rotation_matrix, translation_vector)

            # find the closest points.
            if should_stop.get('last_rms', max_dist) <= 1:
                max_dist = should_stop.get('last_rms') * 3.5

            # translated_cloud, closest_cloud = self._find_correspondences(translated_cloud, target_cloud)
            closest_cloud, translated_cloud = self._find_correspondences(target_cloud, translated_cloud,
                                                                         max_dist=max_dist)

            try:
                if should_stop(translated_cloud, closest_cloud):
                    break
            except RuntimeError as e:
                print(e)
                rotation_matrix = eye(3)
                translation_vector = zeros(3)

            delta_rotation_matrix = self._extract_rotation(translated_cloud, closest_cloud)

            # extract translation.
            source_center = transform_point_cloud(translated_cloud, delta_rotation_matrix).mean(axis=0)
            closest_center = closest_cloud.mean(axis=0)
            delta_translation_vector = round(closest_center - source_center, self._precision)
            # print(">>> delta translation: %s" % str(delta_translation_vector))

            # obtain the translated cloud.
            translated_cloud += delta_translation_vector

            # obtain the transformation (rotation and translation).
            rotation_matrix = matmul(delta_rotation_matrix, rotation_matrix)
            translation_vector += delta_translation_vector

        return rotation_matrix, translation_vector
